{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cdf4490d",
   "metadata": {},
   "source": [
    "# Mutator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5657edf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "\n",
    "import tvm\n",
    "from tvm.relax.frontend import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5264e60d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def test_mutator_naming_basic():\n",
    "    class Module0(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.param0 = nn.Parameter((32, 128), \"float64\")\n",
    "\n",
    "    class Module1(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.mod0 = Module0()\n",
    "            self.param1 = nn.Parameter((32, 128), \"float32\")\n",
    "\n",
    "    class Module2(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.mod1 = Module1()\n",
    "            self.param2 = nn.Parameter((32, 128), \"float16\")\n",
    "\n",
    "    class Module3(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.mod2 = Module2()\n",
    "            self.param3 = nn.Parameter((32, 128), \"float8\")\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_param(self, name: str, node: nn.Parameter) -> Any:\n",
    "            if node.dtype == \"float8\":\n",
    "                assert name == \"mod3.param3\"\n",
    "                return node\n",
    "            elif node.dtype == \"float16\":\n",
    "                assert name == \"mod3.mod2.param2\"\n",
    "                return node\n",
    "            elif node.dtype == \"float32\":\n",
    "                assert name == \"mod3.mod2.mod1.param1\"\n",
    "                return node\n",
    "            elif node.dtype == \"float64\":\n",
    "                assert name == \"mod3.mod2.mod1.mod0.param0\"\n",
    "                return node\n",
    "\n",
    "    mod3 = Module3()\n",
    "    mutator = Mutator()\n",
    "    mutator.visit(\"mod3\", mod3)\n",
    "\n",
    "\n",
    "def test_mutator_naming_modulelist():\n",
    "    class Module(nn.Module):\n",
    "        def __init__(self, dtype) -> None:\n",
    "            super().__init__()\n",
    "            self.param = nn.Parameter((32, 128), dtype)\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_param(self, name: str, node: nn.Parameter) -> Any:\n",
    "            if node.dtype == \"float64\":\n",
    "                assert name == \"mod_list.0.0.param\"\n",
    "                return node\n",
    "            elif node.dtype == \"float32\":\n",
    "                assert name == \"mod_list.0.1.param\"\n",
    "                return node\n",
    "            elif node.dtype == \"float16\":\n",
    "                assert name == \"mod_list.1.0.param\"\n",
    "                return node\n",
    "            elif node.dtype == \"float8\":\n",
    "                assert name == \"mod_list.1.1.param\"\n",
    "                return node\n",
    "\n",
    "    mod_list = nn.ModuleList(\n",
    "        [\n",
    "            nn.ModuleList([Module(\"float64\"), Module(\"float32\")]),\n",
    "            nn.ModuleList([Module(\"float16\"), Module(\"float8\")]),\n",
    "        ]\n",
    "    )\n",
    "    mutator = Mutator()\n",
    "    mutator.visit(\"mod_list\", mod_list)\n",
    "\n",
    "\n",
    "def test_mutator_module():\n",
    "    class SubModule1(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "\n",
    "    class SubModule2(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "\n",
    "    class Module(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.mod = SubModule1()\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_module(self, name: str, node: nn.Module) -> Any:\n",
    "            if isinstance(node, SubModule1):\n",
    "                return SubModule2()\n",
    "            else:\n",
    "                return node\n",
    "\n",
    "    mutator = Mutator()\n",
    "    module = Module()\n",
    "    assert isinstance(module.mod, SubModule1)\n",
    "    module = mutator.visit(\"\", module)\n",
    "    assert isinstance(module.mod, SubModule2)\n",
    "\n",
    "\n",
    "def test_mutator_modulelist():\n",
    "    class Module1(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "\n",
    "    class Module2(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "\n",
    "    class Module3(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_module(self, name: str, node: nn.Module) -> Any:\n",
    "            if isinstance(node, Module3):\n",
    "                return Module1()\n",
    "            else:\n",
    "                return node\n",
    "\n",
    "    mutator = Mutator()\n",
    "    module_list = nn.ModuleList([Module1(), Module2(), Module3()])\n",
    "    assert isinstance(module_list[0], Module1)\n",
    "    assert isinstance(module_list[1], Module2)\n",
    "    assert isinstance(module_list[2], Module3)\n",
    "    module_list = mutator.visit(\"\", module_list)\n",
    "    print(module_list[2])\n",
    "    assert isinstance(module_list[0], Module1)\n",
    "    assert isinstance(module_list[1], Module2)\n",
    "    assert isinstance(module_list[2], Module1)\n",
    "\n",
    "\n",
    "def test_mutator_effect():\n",
    "    class Effect1(nn.Effect):\n",
    "        pass\n",
    "\n",
    "    class Effect2(nn.Effect):\n",
    "        pass\n",
    "\n",
    "    class Module(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.effect = Effect1()\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_effect(self, name: str, node: nn.Effect) -> Any:\n",
    "            if isinstance(node, Effect1):\n",
    "                return Effect2()\n",
    "\n",
    "    mutator = Mutator()\n",
    "    module = Module()\n",
    "    assert isinstance(module.effect, Effect1)\n",
    "    module = mutator.visit(\"\", module)\n",
    "    assert isinstance(module.effect, Effect2)\n",
    "\n",
    "\n",
    "def test_mutator_param():\n",
    "    class Module(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.weight = nn.Parameter((128, 64), \"float16\")\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_param(self, name: str, node: nn.Parameter) -> Any:\n",
    "            if node.dtype == \"float16\":\n",
    "                return nn.Parameter(node.shape, \"float32\")\n",
    "\n",
    "    mutator = Mutator()\n",
    "    module = Module()\n",
    "    assert module.weight.dtype == \"float16\"\n",
    "    module = mutator.visit(\"\", module)\n",
    "    assert module.weight.dtype == \"float32\"\n",
    "\n",
    "\n",
    "def test_mutator_recursively():\n",
    "    class SubModule(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.weight = nn.Parameter((128, 64), \"float16\")\n",
    "\n",
    "    class Module(nn.Module):\n",
    "        def __init__(self) -> None:\n",
    "            super().__init__()\n",
    "            self.mod = SubModule()\n",
    "\n",
    "    class Mutator(nn.Mutator):\n",
    "        def visit_param(self, name: str, node: nn.Parameter) -> Any:\n",
    "            if node.dtype == \"float16\":\n",
    "                return nn.Parameter(node.shape, \"float32\")\n",
    "\n",
    "    mutator = Mutator()\n",
    "    module = Module()\n",
    "    assert module.mod.weight.dtype == \"float16\"\n",
    "    module = mutator.visit(\"\", module)\n",
    "    assert module.mod.weight.dtype == \"float32\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60bb6bca",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
